#!/usr/bin/python3
# -*- coding: utf-8 -*-
import serial
import serial.tools.list_ports

import rospy
from geometry_msgs.msg import Transform,Pose,Point32
from cv_bridge import CvBridge, CvBridgeError

import numpy as np

from bs_img.serial_pack import *

from bs_img.bs_cfg_real_5m import *
 
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import pandas as pd
import traceback


class get_final_res():
    def __init__(self) -> None:
        self.lf_drogue_res_topic = "/bs_debuger/result/lf_drogue"
        self.sf_drogue_res_topic = "/bs_debuger/result/sf_drogue"
        self.wa_drogue_res_topic = "/bs_debuger/result/wa_drogue"
        self.lf_plane_res_topic = "/bs_debuger/result/lf_plane"
        self.sf_plane_res_topic = "/bs_debuger/result/sf_plane"

        self.lf_drogue_res = [np.nan]*3
        self.sf_drogue_res = [np.nan]*3
        self.wa_drogue_res = [np.nan]*3
        self.lf_plane_res = [np.nan]*3
        self.sf_plane_res = [np.nan]*3
        # weight coefficient for the results of lf camera in fusion
        self.coeff = camK_lf[0][0]/(camK_lf[0][0]+camK_sf[0][0])
        self.sleep_interval = 0.0005

        #final result
        self.plane_res = np.array([np.nan, np.nan, np.nan])
        self.drogue_res = np.array([np.nan, np.nan, np.nan])

        self.invalid_frame_num_plane = 0      #compute continus invalid frames
        self.invalid_frame_num_drogue = 0
        self.invalid_frame_max = 5     #if invaild frame < invalid_frame_max, output results of the last frame

        #Kalman Filter
        self.Q = np.eye(3)*0.0025
        self.R = np.eye(3)
        self.x_plane = np.array([np.nan, np.nan, np.nan])
        self.x_drogue = np.array([np.nan, np.nan, np.nan])
        self.P_plane = np.eye(3)
        self.P_drogue = np.eye(3)


        # self.lf_drogue_res_sub = rospy.Subscriber(self.lf_drogue_res_topic, Transform, self.lf_drogue_res_cb)
        # self.sf_drogue_res_sub = rospy.Subscriber(self.sf_drogue_res_topic, Transform, self.sf_drogue_res_cb)
        # self.wa_drogue_res_sub = rospy.Subscriber(self.wa_drogue_res_topic, Transform, self.wa_drogue_res_cb)
        self.lf_plane_res_sub = rospy.Subscriber(self.lf_plane_res_topic, Transform, self.lf_plane_res_cb)
        self.sf_plane_res_sub = rospy.Subscriber(self.sf_plane_res_topic, Transform, self.sf_plane_res_cb)



    def lf_drogue_res_cb(self,transform_msg):
        self.lf_drogue_res = [transform_msg.translation.x, transform_msg.translation.y, transform_msg.translation.z]
        lf_drouge_is_visiable = True if not(np.all(np.isnan(self.lf_drogue_res))) else False
        sf_drouge_is_visiable = True if not(np.all(np.isnan(self.sf_drogue_res))) else False
        wa_drouge_is_visiable = True if not(np.all(np.isnan(self.wa_drogue_res))) else False
        if lf_drouge_is_visiable or sf_drouge_is_visiable or wa_drouge_is_visiable:
            if wa_drouge_is_visiable:
                self.drogue_res = self.wa_drogue_res
            elif lf_drouge_is_visiable and sf_drouge_is_visiable:
                self.drogue_res = np.array(self.lf_drogue_res)*self.coeff + np.array(self.sf_drogue_res)*(1-self.coeff)    
            elif not(lf_drouge_is_visiable) and sf_drouge_is_visiable:
                self.drogue_res = self.sf_drogue_res  
            elif not(sf_drouge_is_visiable) and lf_drouge_is_visiable:
                self.drogue_res = self.lf_drogue_res
            if np.all(np.isnan(self.x_drogue)):   #the first frame
                self.x_drogue = self.drogue_res
            else:
                z_measure_drogue = np.array(self.drogue_res)
                self.P_drogue = self.P_drogue + self.Q
                K = np.matmul(self.P_drogue, np.linalg.pinv(self.P_drogue + self.R))
                self.x_drogue = self.x_drogue + np.matmul(K, z_measure_drogue - self.x_drogue)
                self.P_drogue = np.matmul(np.eye(3) - K, self.P_drogue)  
            drogue_res_frame = serial_frame("drogue",True,round(self.x_drogue[2]*100),round(self.x_drogue[0]*100),round(self.x_drogue[1])*100)    
        else:
            self.invalid_frame_num_drogue = self.invalid_frame_num_drogue + 1
            if self.invalid_frame_num_drogue <= self.invalid_frame_max:
                self.drogue_res = self.x_drogue
                if not(np.all(np.isnan(self.drogue_res))):
                    drogue_res_frame = serial_frame("drogue",True,round(self.x_drogue[2]*100),round(self.x_drogue[0]*100),round(self.x_drogue[1])*100)     
                else:
                    drogue_res_frame = serial_frame("drogue",False,round(0),round(0),round(0))
            else:
                drogue_res_frame = serial_frame("drogue",False,round(0),round(0),round(0))

        uart.write(drogue_res_frame)
        time.sleep(self.sleep_interval)

        
    def sf_drogue_res_cb(self,transform_msg):
        self.sf_drogue_res = [transform_msg.translation.x, transform_msg.translation.y, transform_msg.translation.z]
        lf_drouge_is_visiable = True if not(np.all(np.isnan(self.lf_drogue_res))) else False
        sf_drouge_is_visiable = True if not(np.all(np.isnan(self.sf_drogue_res))) else False
        wa_drouge_is_visiable = True if not(np.all(np.isnan(self.wa_drogue_res))) else False
        if lf_drouge_is_visiable or sf_drouge_is_visiable or wa_drouge_is_visiable:
            if wa_drouge_is_visiable:
                self.drogue_res = self.wa_drogue_res
            elif lf_drouge_is_visiable and sf_drouge_is_visiable:
                self.drogue_res = np.array(self.lf_drogue_res)*self.coeff + np.array(self.sf_drogue_res)*(1-self.coeff)    
            elif not(lf_drouge_is_visiable) and sf_drouge_is_visiable:
                self.drogue_res = self.sf_drogue_res  
            elif not(sf_drouge_is_visiable) and lf_drouge_is_visiable:
                self.drogue_res = self.lf_drogue_res
            if np.all(np.isnan(self.x_drogue)):   #the first frame
                self.x_drogue = self.drogue_res
            else:
                z_measure_drogue = np.array(self.drogue_res)
                self.P_drogue = self.P_drogue + self.Q
                K = np.matmul(self.P_drogue, np.linalg.pinv(self.P_drogue + self.R))
                self.x_drogue = self.x_drogue + np.matmul(K, z_measure_drogue - self.x_drogue)
                self.P_drogue = np.matmul(np.eye(3) - K, self.P_drogue)  
            drogue_res_frame = serial_frame("drogue",True,round(self.x_drogue[2]*100),round(self.x_drogue[0]*100),round(self.x_drogue[1])*100)    
        else:
            self.invalid_frame_num_drogue = self.invalid_frame_num_drogue + 1
            if self.invalid_frame_num_drogue <= self.invalid_frame_max:
                self.drogue_res = self.x_drogue
                if not(np.all(np.isnan(self.drogue_res))):
                    drogue_res_frame = serial_frame("drogue",True,round(self.x_drogue[2]*100),round(self.x_drogue[0]*100),round(self.x_drogue[1])*100)     
                else:
                    drogue_res_frame = serial_frame("drogue",False,round(0),round(0),round(0))
            else:
                drogue_res_frame = serial_frame("drogue",False,round(0),round(0),round(0))

        uart.write(drogue_res_frame)
        time.sleep(self.sleep_interval)

    def wa_drogue_res_cb(self,transform_msg):
        self.wa_drogue_res = [transform_msg.translation.x, transform_msg.translation.y, transform_msg.translation.z]


    def lf_plane_res_cb(self,transform_msg):
        self.lf_plane_res = [transform_msg.translation.x, transform_msg.translation.y, transform_msg.translation.z]
        sf_plane_is_visiable = True if not(np.all(np.isnan(self.sf_plane_res))) else False
        lf_plane_is_visiable = True if not(np.all(np.isnan(self.lf_plane_res))) else False
        if not(sf_plane_is_visiable) and not(lf_plane_is_visiable):
            self.invalid_frame_num_plane = self.invalid_frame_num_plane + 1
            if self.invalid_frame_num_plane <= self.invalid_frame_max:
                self.plane_res = self.x_plane
                if not(np.all(np.isnan(self.plane_res))):
                    plane_res_frame = serial_frame("plane", True, round(self.plane_res[2]*100),round(self.plane_res[0]*100),round(self.plane_res[1]*100))  
                else:
                    plane_res_frame = serial_frame("plane",False,round(0),round(0),round(0))
            else:
                plane_res_frame = serial_frame("plane",False,round(0),round(0),round(0))
        
        else:
            self.invalid_frame_num_plane = 0
            if sf_plane_is_visiable and not(lf_plane_is_visiable):
                self.plane_res = self.sf_plane_res  
            elif lf_plane_is_visiable and not(sf_plane_is_visiable):
                self.plane_res = self.lf_plane_res 
            else:
                self.plane_res = np.array(self.lf_plane_res)*self.coeff + np.array(self.sf_plane_res)*(1-self.coeff)
            if np.all(np.isnan(self.x_plane)):   #the first frame
                self.x_plane = self.plane_res
            else:
                z_measure_plane = np.array(self.plane_res)
                self.P_plane = self.P_plane + self.Q
                K = np.matmul(self.P_plane, np.linalg.pinv(self.P_plane + self.R))
                self.x_plane = self.x_plane + np.matmul(K, z_measure_plane - self.x_plane)
                self.P_plane = np.matmul(np.eye(3) - K, self.P_plane)
            plane_res_frame = serial_frame("plane",True,round(self.x_plane[2]*100),round(self.x_plane[0]*100),round(self.x_plane[1]*100))  
        

        uart.write(plane_res_frame)
        time.sleep(self.sleep_interval)

        print(self.x_plane)


    def sf_plane_res_cb(self,transform_msg):
        self.sf_plane_res = [transform_msg.translation.x, transform_msg.translation.y, transform_msg.translation.z]
        sf_plane_is_visiable = True if not(np.all(np.isnan(self.sf_plane_res))) else False
        lf_plane_is_visiable = True if not(np.all(np.isnan(self.lf_plane_res))) else False
        if not(sf_plane_is_visiable) and not(lf_plane_is_visiable):
            self.invalid_frame_num_plane = self.invalid_frame_num_plane + 1
            if self.invalid_frame_num_plane <= self.invalid_frame_max:
                self.plane_res = self.x_plane
                if not(np.all(np.isnan(self.plane_res))):
                    plane_res_frame = serial_frame("plane", True, round(self.plane_res[2]*100),round(self.plane_res[0]*100),round(self.plane_res[1]*100))  
                else:
                    plane_res_frame = serial_frame("plane",False,round(0),round(0),round(0))
            else:
                plane_res_frame = serial_frame("plane",False,round(0),round(0),round(0))
        
        else:
            self.invalid_frame_num_plane = 0
            if sf_plane_is_visiable and not(lf_plane_is_visiable):
                self.plane_res = self.sf_plane_res  
            elif lf_plane_is_visiable and not(sf_plane_is_visiable):
                self.plane_res = self.lf_plane_res 
            else:
                self.plane_res = np.array(self.lf_plane_res)*self.coeff + np.array(self.sf_plane_res)*(1-self.coeff)
            if np.all(np.isnan(self.x_plane)):   #the first frame
                self.x_plane = self.plane_res
            else:
                z_measure_plane = np.array(self.plane_res)
                self.P_plane = self.P_plane + self.Q
                K = np.matmul(self.P_plane, np.linalg.pinv(self.P_plane + self.R))
                self.x_plane = self.x_plane + np.matmul(K, z_measure_plane - self.x_plane)
                self.P_plane = np.matmul(np.eye(3) - K, self.P_plane)
            plane_res_frame = serial_frame("plane",True,round(self.x_plane[2]*100),round(self.x_plane[0]*100),round(self.x_plane[1]*100))  
        

        uart.write(plane_res_frame)
        time.sleep(self.sleep_interval)

        print(self.x_plane)
        



if __name__ == "__main__":
    rospy.init_node('bs_solver_node', anonymous=True)

    try:
        port, baud =  '/dev/ttyUSB0', 115200
        uart = serial.Serial(port,baud,parity=serial.PARITY_NONE)

        final_res = get_final_res()
    except Exception as e:
        print(traceback.format_exc())
    rospy.spin()